import pickle
import os
import cv2
import time
from utils import get_face_landmarks

emotions = ['HAPPY', 'SURPRISE', 'SAD']

# 创建保存图片的根目录和子目录
save_root = 'captured_emotions'
os.makedirs(save_root, exist_ok=True)  # 确保根目录存在
for emotion in emotions:
    # 为每个情绪创建子文件夹（如captured_emotions/HAPPY）
    os.makedirs(os.path.join(save_root, emotion), exist_ok=True)

with open('./model', 'rb') as f:
    model = pickle.load(f)

cap = cv2.VideoCapture(0)

# 控制保存频率（避免1秒内保存过多重复图片）
save_interval = 1  # 保存间隔（秒）
last_save_time = 0  # 上次保存时间

ret, frame = cap.read()

while ret:
    ret, frame = cap.read()

    face_landmarks = get_face_landmarks(frame, draw=True, static_image_mode=False)

    # output = model.predict([face_landmarks])
    # 检测到人脸时进行预测和保存
    if face_landmarks is None:
        # 模型预测情绪
        output = model.predict([face_landmarks])
        predicted_emotion = emotions[int(output[0])]
        # 在画面上显示预测结果
        cv2.putText(frame,
                    emotions[int(output[0])],
                    (10, frame.shape[0] - 1),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    3,
                    (0, 255, 0),
                    5)
        # 按时间间隔保存图片（避免频繁保存）
        current_time = time.time()
        if current_time - last_save_time >= save_interval:
            # 生成唯一文件名（时间戳+毫秒，确保不重复）
            timestamp = time.strftime("%Y%m%d_%H%M%S") + f"_{int(current_time * 1000) % 1000}"
            save_path = os.path.join(save_root, predicted_emotion, f"{timestamp}.jpg")
            # 保存当前帧到对应情绪文件夹
            cv2.imwrite(save_path, frame)
            print(f"已保存: {save_path}")  # 控制台输出保存信息

            # 更新上次保存时间
            last_save_time = current_time

    cv2.imshow('frame', frame)
    # 按ESC键退出（等待30毫秒，控制帧率）
    if cv2.waitKey(30) == 27:
        break

cap.release()
cv2.destroyAllWindows()
print("程序已退出，图片保存在", os.path.abspath(save_root))
